import {Client} from "@modelcontextprotocol/sdk/client/index.js";
import {SSEClientTransport} from "@modelcontextprotocol/sdk/client/sse.js";
import dotenv from "dotenv";
import OpenAI from "openai";
import readline from "readline/promises";

// 加载环境变量
dotenv.config();
const LLM_API_KEY = process.env.LLM_API_KEY;
const LLM_BASE_URL = process.env.LLM_BASE_URL;
const LLM_MODEL = process.env.LLM_MODEL;

// 创建客户端实例
const client = new Client({
    name: "example-client",
    version: "1.0.0"
}, {
    capabilities: {}
});

// 初始化 sse 传输实例
const transport = new SSEClientTransport(
    new URL("http://localhost:3000/sse")
);

// 连接到服务器
await client.connect(transport);

// 获取工具列表
const toolsResult = await client.listTools();
// 处理工具列表，转换为特定格式
const tools: OpenAI.ChatCompletionTool[] = toolsResult.tools.map((tool) => {
    return {
        type: "function",
        function: {
            name: tool.name,
            description: tool.description,
            parameters: tool.inputSchema,
        }
    };
});
// 打印连接成功信息和工具名称
console.log(
    "成功链接服务器，获取到工具如下:",
    tools.map((item) => item.function.name)
);

// 创建 readline 接口
const rl = readline.createInterface({
    input: process.stdin,
    output: process.stdout,
});
// 获取用户输入
const query = await rl.question("\n请输入: ");

// 初始化 openai 实例，传入 API 密钥和基础 URL
const openai = new OpenAI({
    apiKey: LLM_API_KEY,
    baseURL: LLM_BASE_URL
});
// 初始化消息数组
const messages: OpenAI.ChatCompletionMessageParam[] = [
    {
        role: "user",
        content: query,
    },
];
// 调用 OpenAI 的聊天完成接口
const response = await openai.chat.completions.create({
    model: LLM_MODEL!,
    max_tokens: 1000,
    messages,
    tools: tools,
});
// openai 非流式对话，只有一个 choice
const resMessage = response.choices[0].message
// 如果没有工具调用
if (!resMessage.tool_calls || resMessage.tool_calls.length === 0) {
    // 直接输出回答
    console.log(resMessage.content)
} else {
    // 大模型回答加入上下文
    messages.push(resMessage)

    // 获取工具ID
    const toolId = resMessage.tool_calls[0].id
    // 获取工具名称
    const toolName = resMessage.tool_calls[0].function.name;
    // 获取工具参数
    const toolArgs = JSON.parse(resMessage.tool_calls[0].function.arguments);

    // 调用工具
    const toolResult = await client.callTool({
        name: toolName,
        arguments: toolArgs,
    });
    // 将工具调用信息添加到最终
    console.log(`\n[调用工具 ${toolName} 参数为 ${JSON.stringify(toolArgs)} 结果为 ${JSON.stringify(toolResult)}]\n`)

    // 将工具调用结果添加到消息数组
    messages.push({
        role: "tool",
        content: JSON.stringify(toolResult.content),
        tool_call_id: toolId
    });

    // 再次调用 OpenAI 的聊天完成接口
    const response = await openai.chat.completions.create({
        model: LLM_MODEL!,
        max_tokens: 1000,
        // 此时消息列表中有三条消息：1用户提问 2大模型返回工具调用 3工具调用结果
        messages,
    });
    console.log(response.choices[0].message.content)

}

// 关闭客户端
client.close()
process.exit(0);